Introduce differentiation interface #2137

SouthEndMusic · 2025-03-10T10:53:27Z

SouthEndMusic · 2025-03-10T10:57:20Z

@gdalle it would be nice if you could have a look at 64091bc. This already works very nicely! We might be able to get rid of our direct SparseConnectivityTracer dependency, were it not that we (I) did some naughty overloads using GradientTracer. What would the next step towards Enzyme support look like?

SouthEndMusic · 2025-03-10T13:37:32Z

When testing dense Jacobian + ForwardDiff I get an error like this:

ERROR: LoadError: Invalid Tag object: ...
Observed ForwardDiff.Tag{...}

I wonder whether that has something to do with passing the AD backend to both DifferentiationInterface.jacobian and the algorithm constructor. I also wonder how this works for algorithms that require derivatives .w.r.t. time.

gdalle · 2025-03-10T13:41:40Z

Taking a look now

Manifest.toml

Project.toml

core/src/Ribasim.jl

core/src/model.jl

gdalle · 2025-03-10T13:46:19Z

core/src/model.jl

@@ -28,6 +28,33 @@ struct Model{T}
    end
 end

+function get_jac_eval(du::Vector, u::Vector, p::Parameters, solver::Solver)
+    backend = if solver.autodiff
+        AutoForwardDiff()


Maybe you want to configure the tag here if it is available from somewhere else (perhaps the solver?)

solver is our own solver config object, but it works if I consistently specify the same tag everywhere 👍

gdalle · 2025-03-10T13:48:22Z

core/src/util.jl

-    p.all_nodes_active = false
-    jac_prototype
-end
-
 # Custom overloads


If you want to get rid of the SCT dependency, you may want to

put those overloads in an SCT package extension

ask the user to provide an AD backend, and if it is an AutoSparse, retrieve its sparsity_detector instead of providing your own

Hmm this probably doesn't fit our application as described in #2137 (comment). Keeping the SCT dependency is fine

gdalle · 2025-03-10T13:48:58Z

core/src/util.jl

 # Custom overloads
-reduction_factor(x::GradientTracer, threshold::Real) = x
+reduction_factor(x::GradientTracer, ::Real) = x


Note that this is explicitly discouraged in the SCT docs: see the following page to add overloads properly

https://adrianhill.de/SparseConnectivityTracer.jl/stable/internals/adding_overloads/

I know I know 🙃

Some of these functions have more than 2 arguments, but for all of them we only care about the derivative with respect to one input. It's not clear to me whether/how that fits within the overload functionality

It makes sense to use our mechanisms. They will generate a couple of superfluous methods, but I don't see the harm in that.

To be fair, this specific line of code looks harmless. But once you have more than a handful of overloads, you might want to create an SCT extension and follow our docs. The generated code will be more future proof and compatible with local and global Jacobian and Hessian tracers. Your current code only supports global Jacobians.

core/src/Ribasim.jl

gdalle · 2025-03-10T13:51:36Z

core/src/model.jl

+
+    # Activate all nodes to catch all possible state dependencies
+    p.all_nodes_active = true
+    prep = prepare_jacobian((du, u) -> water_balance!(du, u, p, t), du, backend, u)


Suggested change

prep = prepare_jacobian((du, u) -> water_balance!(du, u, p, t), du, backend, u)

prep = prepare_jacobian(water_balance!, du, backend, u, Constant(p), Constant(t))

See https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/advanced/#Contexts

From the docs:

Another option would be creating a closure, but that is sometimes undesirable.

When is a closure undesirable?

gradient(f, backend, x, Constant(c))
gradient(f, backend, x, Cache(c))

In the first call, c is kept unchanged throughout the function evaluation. In the second call, c can be mutated with values computed during the function.

Our p contains caches for in place computations in our RHS (hence the discussion on PreallocationTools etc. in the related issue). does that mean that we should use Cache(p)?

Should SciMLStructures.jl come in here for more granular control?

When is a closure undesirable?

With Enzyme in particular it can make things slower or even error. With other backends it doesn't make much of a difference, but explicitly laying out the Contexts also allows taking into account element types (eg. for handling translation to Dual with Caches).

Our p contains caches for in place computations in our RHS (hence the discussion on PreallocationTools etc. in the related issue). does that mean that we should use Cache(p)?

Does p contain anything whose value you care about?
In general, you might want to split it between a Constant part and a Cache part.

Should SciMLStructures.jl come in here for more granular control?

DI has no specific support for SciMLStructures

@gdalle I'm working on a refactor where e.g. the prep definition now looks like this:

prep = prepare_jacobian( water_balance!, du, backend, u, Constant(p_non_diff), Cache(p_diff), Constant(p_mutable), Constant(t), )

This now fails in the sparsity detection because of an attempt to write GradientTracer values to a Vector{Float64} field of p_diff::ParametersDiff. I made ParametersDiff parametric so ParametersDiff{GradientTracer{...}} can exist, and I half expected this to be constructed internally. This probably worked before because of PreallocationTools.

Oh I just saw this warning in the docs:

Most backends require any Cache context to be an AbstractArray.

Let's see what I can do with that.

It doesn't look like that quickly solves the problem. I just naively subtyped ParametersDiff{T} <: AbstractVector{T}. Maybe I need to overload some methods?

MWE: JuliaDiff/DifferentiationInterface.jl#737

Thanks, that's a tricky one and it's indeed on me. If you don't want to wait for a DI fix (ETA ~ days), a short-term solution would be to use PreallocationTools and a closure, even if it makes Enzyme angry.

gdalle · 2025-03-10T13:52:14Z

core/src/model.jl

+    jac =
+        (J, u, p, t) ->
+            jacobian!((du, u) -> water_balance!(du, u, p, t), du, J, prep, backend, u)


Suggested change

jac =

(J, u, p, t) ->

jacobian!((du, u) -> water_balance!(du, u, p, t), du, J, prep, backend, u)

jac(J, u, p, t) = jacobian!(water_balance!, du, J, prep, backend, u, Constant(p), Constant(t))

gdalle · 2025-03-13T07:14:17Z

Pass most tests

Love the confident commit naming. That's the spirit we wanna see

gdalle · 2025-03-13T07:20:35Z

core/src/model.jl

+    diff_cache_SCT =
+        zeros(GradientTracer{IndexSetGradientPattern{Int64, BitSet}}, length(diff_cache))


With this PR you can use SCT.jacobian_buffer instead, and with the update to DI I'll make once that is merged, you probably won't need any tweak at all

gdalle · 2025-03-13T16:48:28Z

@SouthEndMusic can you take the branch from JuliaDiff/DifferentiationInterface.jl#739 for a spin, see if it works?

Warning

I am aware that using DI.Cache still gives rise to allocations during each Jacobian computation. You're the first person actually using it, so I plan to fix it, but first I want to know if it runs

SouthEndMusic · 2025-03-13T17:59:25Z

@gdalle I took the main branch, and it indeed works 👍

gdalle · 2025-03-14T13:10:25Z

@SouthEndMusic with the branch from JuliaDiff/DifferentiationInterface.jl#741 the Caches should be allocation-free after preparation.

gdalle · 2025-03-15T11:36:18Z

The changes have been released

SouthEndMusic · 2025-03-18T15:37:47Z

Some runtimes of the \hws_2024_7_0 model:

# FiniteDiff: 22.592886
# ForwardDiff: 20.455309
# ForwardDiff + type instability fix: 17.090883

gdalle · 2025-03-18T15:40:50Z

If the Jacobian is sparse, can you try other orders inside the GreedyColoringAlgorithm? For instance, GreedyColoringAlgorithm(LargestFirst()) or GreedyColoringAlgorithm(RandomOrder(rng, seed))?

SouthEndMusic · 2025-03-18T15:44:48Z

If the Jacobian is sparse, can you try other orders inside the GreedyColoringAlgorithm? For instance, GreedyColoringAlgorithm(LargestFirst()) or GreedyColoringAlgorithm(RandomOrder(rng, seed))?

What can be the effect of this? A different number of rhs calls required to compute the Jacobian?

gdalle · 2025-03-18T15:47:38Z

Yes, it influences the number of different colors with which the columns of the Jacobian are colored, and one color equals one function call (not exactly true with ForwardDiff though). This could accelerate the FiniteDiff version significantly if the natural coloring was suboptimal.

gdalle · 2025-03-18T15:50:07Z

You can call SparseMatrixColorings.ncolors(prep) on the Jacobian preparation result to see how many colors you have. For a dense Jacobian, this is equal to the total number of columns (input dimension). And the rule is that the cost of the sparse Jacobian scales with the number of colors (or the number of colors divided by 12 for ForwardDiff).

SouthEndMusic · 2025-03-18T15:52:36Z

LargestFirst() doesn't seem to have a significant effect

gdalle · 2025-03-18T18:10:52Z

Heads up, DI v0.6.46 supports nested tuples and named tuples of arrays as caches, at least for the most common backends

SouthEndMusic · 2025-03-19T10:28:52Z

Heads up, DI v0.6.46 supports nested tuples and named tuples of arrays as caches, at least for the most common backends

This is really nice, quite a bit of complexity can be removed now 👍

visr

Great stuff @SouthEndMusic. And thanks for all the help @gdalle.

I left some minor comments, but let's get this in soon, we can always refine later.

core/src/parameter.jl

core/src/util.jl

visr · 2025-03-21T10:33:24Z

core/src/validation.jl

        @assert flow_rate_update.name == :flow_rate
        flow_rate_ = minimum(flow_rate_update.value.u)

        if flow_rate_ < 0.0
            errors = true
            control_state = key[2]
-            @error "$id_controlled flow rates must be non-negative, found $flow_rate_ for control state '$control_state'."
+            @error "Negative flow rate(s) for $id_controlled, control state '$control_state' found."


Suggested change

@error "Negative flow rate(s) for $id_controlled, control state '$control_state' found."

@error "Negative flow rate(s) found." node_id=id_controlled control_state

@gijsber

This adds a function `model.to_fews(region_dir)` that converts the network and results to files that Delft-FEWS can directly handle. It is marked as experimental for now. @gijsber is working on a Delft-FEWS configuration that can be used to visualize model results, to complement our existing tools. We'll likely add this configuration to this monorepo since it is generic. #2159 also pertains to this work. What is especially nice is the spatio-temporal support of Delft-FEWS, so we can make visualizations like this: ![image](https://github.com/user-attachments/assets/2e61bf82-0d7d-4558-a645-755d7e763b74) In theory we can support similar functionality with QGIS, but looking at the plots in #1369 this would likely need work in QGIS itself. So this is really a quick win to be able to inspect models better. --------- Co-authored-by: Maarten Pronk <git@evetion.nl>

SouthEndMusic added 2 commits March 10, 2025 10:22

Unrelated cleanups

d3c68c2

Introduce DifferentiationInterface

64091bc

Fix most tests

732f074

bring back TracerSparsityDetector

b5a5693

gdalle suggested changes Mar 10, 2025

View reviewed changes

Most comments adressed

387d3d4

SouthEndMusic marked this pull request as draft March 10, 2025 15:58

SouthEndMusic added 6 commits March 11, 2025 14:49

Subdividing parameters first part

46af5f4

Merge branch 'main' into introduce_differentiation_interface

e5540f6

Introduce vector cache

79162e2

Add SCT cache hack, fix minor bugs

448ca28

Fix many tests

828178a

Pass most tests

de0dbb7

gdalle reviewed Mar 13, 2025

View reviewed changes

SouthEndMusic added 5 commits March 13, 2025 15:56

Fix more tests

17e2da8

Merge branch 'main' into introduce_differentiation_interface

0c7222c

Replace weight function by Returns

eb31535

Replace weight function by Returns properly

1939178

Avoid accessing undefined cache entries

23cf2cd

We no longer need to explicitly create the Jacobian buffer for SCT

151c552

Use DI prep_cache branch

23b844b

Update DI

a77042e

SouthEndMusic added 2 commits March 18, 2025 10:54

Small docs fix

93aa7a2

Merge branch 'main' into introduce_differentiation_interface

0f9479f

gdalle mentioned this pull request Mar 18, 2025

Support more complex caches JuliaDiff/DifferentiationInterface.jl#747

Closed

SouthEndMusic added 3 commits March 18, 2025 12:37

Actually fix the call stacks in the docs this time

f96feb5

Fix validation docs

2a6df87

Merge branch 'main' into introduce_differentiation_interface

b11ff11

Fix discrete control table + type instability

46735b0

SouthEndMusic added 2 commits March 19, 2025 10:17

Merge branch 'main' into introduce_differentiation_interface

7fb12f9

Fix allocation problem docs example

b550ed2

SouthEndMusic marked this pull request as draft March 19, 2025 09:32

Use NamedTuple DiffCache

78e0497

SouthEndMusic marked this pull request as ready for review March 19, 2025 10:56

Use latest version of the new dependencies

8735c4e

visr approved these changes Mar 21, 2025

View reviewed changes

visr mentioned this pull request Mar 21, 2025

Implement SciMLStructures interface on Parameters #1174

Open

visr and others added 5 commits March 21, 2025 13:00

Don't recompute low storage factor for basin forcings

1d1b97b

Merge branch 'main' into introduce_differentiation_interface

41f0374

Update docstrings

fda4501

Fancify error message

49334b5

SouthEndMusic merged commit a76ab8b into main Mar 25, 2025
19 of 20 checks passed

SouthEndMusic deleted the introduce_differentiation_interface branch March 25, 2025 12:50

	prep = prepare_jacobian((du, u) -> water_balance!(du, u, p, t), du, backend, u)
	prep = prepare_jacobian(water_balance!, du, backend, u, Constant(p), Constant(t))

		diff_cache_SCT =
		zeros(GradientTracer{IndexSetGradientPattern{Int64, BitSet}}, length(diff_cache))

	@error "Negative flow rate(s) for $id_controlled, control state '$control_state' found."
	@error "Negative flow rate(s) found." node_id=id_controlled control_state

Introduce differentiation interface #2137

Introduce differentiation interface #2137

Conversation

SouthEndMusic commented Mar 10, 2025

SouthEndMusic commented Mar 10, 2025

SouthEndMusic commented Mar 10, 2025

gdalle commented Mar 10, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gdalle Mar 10, 2025 • edited Loading

Choose a reason for hiding this comment

SouthEndMusic Mar 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SouthEndMusic Mar 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gdalle commented Mar 13, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gdalle commented Mar 13, 2025

SouthEndMusic commented Mar 13, 2025

gdalle commented Mar 14, 2025

gdalle commented Mar 15, 2025

SouthEndMusic commented Mar 18, 2025

gdalle commented Mar 18, 2025

SouthEndMusic commented Mar 18, 2025

gdalle commented Mar 18, 2025

gdalle commented Mar 18, 2025

SouthEndMusic commented Mar 18, 2025

gdalle commented Mar 18, 2025

SouthEndMusic commented Mar 19, 2025

visr left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gdalle Mar 10, 2025 •

edited

Loading

SouthEndMusic Mar 11, 2025 •

edited

Loading

SouthEndMusic Mar 11, 2025 •

edited

Loading